package top.continew.ai.config;

import com.alibaba.cloud.ai.memory.redis.JedisRedisChatMemoryRepository;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.math3.analysis.function.Abs;
import org.apache.poi.ss.formula.functions.T;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;
import top.continew.ai.advisor.LogAdvisor;
import top.continew.ai.advisor.ReReadingAdvisor;
import top.continew.ai.dto.AiProperties;

import java.util.Objects;

/**
 * chat工具的主入口
 *
 * @Author： Lonni
 * @Date： 2025/8/21  19:09
 */
@Slf4j
@Component
public class ChatApp {

    private final ChatClient chatClient;
    private final AiProperties properties;

    /**
     * AI 调用工具能力
     * 调用本地工具
     */
    @Resource
    private ToolCallback[] allTools;



    public ChatApp(ChatClient.Builder builder, AiProperties properties, JedisRedisChatMemoryRepository chatMemoryRepository) {
        this.properties = properties;
        this.chatClient = builder
                .defaultSystem(properties.getSystemPrompt())
                .defaultAdvisors(
                        //增加多伦会话记忆功能,使用jedis存储
                        MessageChatMemoryAdvisor.builder(
                                MessageWindowChatMemory.builder().
                                        chatMemoryRepository(chatMemoryRepository)
                                        .maxMessages(properties.getMemoryCount())
                                        .build()).build(),
                        //增加日志打印
                        new LogAdvisor(),
                        //增加重读
                        new ReReadingAdvisor()
                )


                .build();
    }


    /**
     * 简单的同步会话功能
     *
     * @param message        用户消息
     * @param conversationId 会话id
     * @return
     */
    public String doChat(String message, String conversationId) {
        ChatResponse chatResponse = this.chatClient
                .prompt()
                .user(message)
                .advisors(p -> p.param(ChatMemory.CONVERSATION_ID, conversationId))
                .call()
                .chatResponse();

        String text = chatResponse.getResult().getOutput().getText();
        log.info("doChat---content--->{}", text);
        return text;
    }

    /**
     * 简单的同步会话功能  sse方式
     *
     * @param message        用户消息
     * @param conversationId 会话id
     * @return
     */
    public Flux<String> doChatByStream(String message, String conversationId) {
        Flux<String> content = this.chatClient
                .prompt()
                .user(message)
                .advisors(p -> p.param(ChatMemory.CONVERSATION_ID, conversationId))
               //增加工具调用
                .toolCallbacks(allTools)
                .stream().content();

//  订阅消息       content.subscribe(text -> log.info("doChat---content--->{}", text));


        return content;
    }



    @Resource
    VectorStore vectorStore;
    @Resource
    RewriteQueryTransformer rewriteQueryTransformer;
    @Resource
    ContextualQueryAugmenter contextualQueryAugmenter;


    /**
     * 增强型的RAG检索
     * @param message 消息
     * @param conversationId  会话id
     * @param topK 返回的文档数  方法参数>配置参数
     * @param expression 过滤条件,如:
     *                   new FilterExpressionBuilder()
     *                 .eq("status", status)
     *                 .build();
     * @return
     */
    public Flux<String> doChatWithRag(String message,
                                      String conversationId,
                                      Integer topK,
                                      Filter.Expression expression
                                      ) {
        Flux<String> content = this.chatClient.prompt()
                .user(message)
                .advisors(p -> p.param(ChatMemory.CONVERSATION_ID, conversationId))
                .advisors(new LogAdvisor())
                .advisors(RetrievalAugmentationAdvisor.builder()
                        //配置基于向量库的查询
                        .documentRetriever(VectorStoreDocumentRetriever.builder()
                                .vectorStore(vectorStore)
                                .topK(Objects.isNull(topK)? properties.getTopK() : topK)
                                .similarityThreshold(properties.getSimilarityThreshold())
                                .filterExpression( expression)
                                .build())
                        //配置空查询时返回的提示
                        .queryAugmenter(contextualQueryAugmenter)
                        //重写查询
                        .queryTransformers(rewriteQueryTransformer)
                        //配置后置监控
                        .documentPostProcessors(((query, documents) -> {
                            log.info("Original query:{} ", query.text());
                            log.info("Retrieved documents: {}", documents.size());
                            return documents;
                        }))
                        .build())
                //增加工具调用
                .toolCallbacks(allTools)
                .stream().content();
        return content;

    }


}
